// https://github.com/8891689
// zuc.c
#include "zuc.h"
#include <string.h>
#include <stdint.h>

/* 内部状态 */
static uint32_t lfsr[16];   // 31-bit LFSR
static uint32_t R1, R2;     // F函数寄存器
static uint32_t X0, X1, X2, X3; // 位重组结果

/* S-Box 和 D 常量  */

static const uint8_t S0[256] =  {
0x3e,0x72,0x5b,0x47,0xca,0xe0,0x00,0x33,0x04,0xd1,0x54,0x98,0x09,0xb9,0x6d,0xcb,
0x7b,0x1b,0xf9,0x32,0xaf,0x9d,0x6a,0xa5,0xb8,0x2d,0xfc,0x1d,0x08,0x53,0x03,0x90,
0x4d,0x4e,0x84,0x99,0xe4,0xce,0xd9,0x91,0xdd,0xb6,0x85,0x48,0x8b,0x29,0x6e,0xac,
0xcd,0xc1,0xf8,0x1e,0x73,0x43,0x69,0xc6,0xb5,0xbd,0xfd,0x39,0x63,0x20,0xd4,0x38,
0x76,0x7d,0xb2,0xa7,0xcf,0xed,0x57,0xc5,0xf3,0x2c,0xbb,0x14,0x21,0x06,0x55,0x9b,
0xe3,0xef,0x5e,0x31,0x4f,0x7f,0x5a,0xa4,0x0d,0x82,0x51,0x49,0x5f,0xba,0x58,0x1c,
0x4a,0x16,0xd5,0x17,0xa8,0x92,0x24,0x1f,0x8c,0xff,0xd8,0xae,0x2e,0x01,0xd3,0xad,
0x3b,0x4b,0xda,0x46,0xeb,0xc9,0xde,0x9a,0x8f,0x87,0xd7,0x3a,0x80,0x6f,0x2f,0xc8,
0xb1,0xb4,0x37,0xf7,0x0a,0x22,0x13,0x28,0x7c,0xcc,0x3c,0x89,0xc7,0xc3,0x96,0x56,
0x07,0xbf,0x7e,0xf0,0x0b,0x2b,0x97,0x52,0x35,0x41,0x79,0x61,0xa6,0x4c,0x10,0xfe,
0xbc,0x26,0x95,0x88,0x8a,0xb0,0xa3,0xfb,0xc0,0x18,0x94,0xf2,0xe1,0xe5,0xe9,0x5d,
0xd0,0xdc,0x11,0x66,0x64,0x5c,0xec,0x59,0x42,0x75,0x12,0xf5,0x74,0x9c,0xaa,0x23,
0x0e,0x86,0xab,0xbe,0x2a,0x02,0xe7,0x67,0xe6,0x44,0xa2,0x6c,0xc2,0x93,0x9f,0xf1,
0xf6,0xfa,0x36,0xd2,0x50,0x68,0x9e,0x62,0x71,0x15,0x3d,0xd6,0x40,0xc4,0xe2,0x0f,
0x8e,0x83,0x77,0x6b,0x25,0x05,0x3f,0x0c,0x30,0xea,0x70,0xb7,0xa1,0xe8,0xa9,0x65,
0x8d,0x27,0x1a,0xdb,0x81,0xb3,0xa0,0xf4,0x45,0x7a,0x19,0xdf,0xee,0x78,0x34,0x60
};

static const uint8_t S1[256] = {
0x55,0xc2,0x63,0x71,0x3b,0xc8,0x47,0x86,0x9f,0x3c,0xda,0x5b,0x29,0xaa,0xfd,0x77,
0x8c,0xc5,0x94,0x0c,0xa6,0x1a,0x13,0x00,0xe3,0xa8,0x16,0x72,0x40,0xf9,0xf8,0x42,
0x44,0x26,0x68,0x96,0x81,0xd9,0x45,0x3e,0x10,0x76,0xc6,0xa7,0x8b,0x39,0x43,0xe1,
0x3a,0xb5,0x56,0x2a,0xc0,0x6d,0xb3,0x05,0x22,0x66,0xbf,0xdc,0x0b,0xfa,0x62,0x48,
0xdd,0x20,0x11,0x06,0x36,0xc9,0xc1,0xcf,0xf6,0x27,0x52,0xbb,0x69,0xf5,0xd4,0x87,
0x7f,0x84,0x4c,0xd2,0x9c,0x57,0xa4,0xbc,0x4f,0x9a,0xdf,0xfe,0xd6,0x8d,0x7a,0xeb,
0x2b,0x53,0xd8,0x5c,0xa1,0x14,0x17,0xfb,0x23,0xd5,0x7d,0x30,0x67,0x73,0x08,0x09,
0xee,0xb7,0x70,0x3f,0x61,0xb2,0x19,0x8e,0x4e,0xe5,0x4b,0x93,0x8f,0x5d,0xdb,0xa9,
0xad,0xf1,0xae,0x2e,0xcb,0x0d,0xfc,0xf4,0x2d,0x46,0x6e,0x1d,0x97,0xe8,0xd1,0xe9,
0x4d,0x37,0xa5,0x75,0x5e,0x83,0x9e,0xab,0x82,0x9d,0xb9,0x1c,0xe0,0xcd,0x49,0x89,
0x01,0xb6,0xbd,0x58,0x24,0xa2,0x5f,0x38,0x78,0x99,0x15,0x90,0x50,0xb8,0x95,0xe4,
0xd0,0x91,0xc7,0xce,0xed,0x0f,0xb4,0x6f,0xa0,0xcc,0xf0,0x02,0x4a,0x79,0xc3,0xde,
0xa3,0xef,0xea,0x51,0xe6,0x6b,0x18,0xec,0x1b,0x2c,0x80,0xf7,0x74,0xe7,0xff,0x21,
0x5a,0x6a,0x54,0x1e,0x41,0x31,0x92,0x35,0xc4,0x33,0x07,0x0a,0xba,0x7e,0x0e,0x34,
0x88,0xb1,0x98,0x7c,0xf3,0x3d,0x60,0x6c,0x7b,0xca,0xd3,0x1f,0x32,0x65,0x04,0x28,
0x64,0xbe,0x85,0x9b,0x2f,0x59,0x8a,0xd7,0xb0,0x25,0xac,0xaf,0x12,0x03,0xe2,0xf2
};

/* D常量 (15位值) */
static const uint16_t D[16] = {
    0x44D7, 0x26BC, 0x626B, 0x135E, 0x5789, 0x35E2, 0x7135, 0x09AF,
    0x4D78, 0x2F13, 0x6BC4, 0x1AF1, 0x5E26, 0x3C4D, 0x789A, 0x47AC
};

/* 31-bit 循環左移 */
static inline uint32_t rotl31(uint32_t x, int n) {
    n %= 31;
    return ((x << n) | (x >> (31 - n))) & 0x7FFFFFFF;
}


/* 高性能但LFSR上下文中适用 */
static inline uint32_t mod_add31(uint32_t a, uint32_t b) {
    uint32_t sum = a + b;
    return (sum & 0x7FFFFFFF) + (sum >> 31);
}

/* LFSR 初始化模式 - 修正嵌套問題 */
static void lfsr_init_mode(uint32_t u) {
    uint32_t v = 0;

    v = mod_add31(v, rotl31(lfsr[15], 15));
    v = mod_add31(v, rotl31(lfsr[13], 17));
    v = mod_add31(v, rotl31(lfsr[10], 21));
    v = mod_add31(v, rotl31(lfsr[4], 20));
    v = mod_add31(v, rotl31(lfsr[0], 8));
    v = mod_add31(v, lfsr[0]);

    v = mod_add31(v, u); 

    if (v == 0) {
        v = 0x7FFFFFFF;
    }

    for (int i = 0; i < 15; i++) {
        lfsr[i] = lfsr[i+1];
    }
    lfsr[15] = v;
}

/* LFSR 工作模式 - 修正嵌套問題 */
static void lfsr_work_mode() {
    uint32_t v = 0;
    v = mod_add31(v, rotl31(lfsr[15], 15));
    v = mod_add31(v, rotl31(lfsr[13], 17));
    v = mod_add31(v, rotl31(lfsr[10], 21));
    v = mod_add31(v, rotl31(lfsr[4], 20));
    v = mod_add31(v, rotl31(lfsr[0], 8));
    v = mod_add31(v, lfsr[0]);

    if (v == 0) {
        v = 0x7FFFFFFF;
    }

    for (int i = 0; i < 15; i++) {
        lfsr[i] = lfsr[i+1];
    }
    lfsr[15] = v;
}

/* 位重組 */
static void bit_reorganize() {
    // 使用移位操作確保正確的位位置
    X0 = ((lfsr[15] >> 15) << 16) | (lfsr[14] & 0xFFFF);
    X1 = ((lfsr[11] & 0xFFFF) << 16) | (lfsr[9] >> 15);
    X2 = ((lfsr[7] & 0xFFFF) << 16) | (lfsr[5] >> 15);
    X3 = ((lfsr[2] & 0xFFFF) << 16) | (lfsr[0] >> 15);
}

/* 線性變換 L1 */
static uint32_t L1(uint32_t x) {
    // 使用32位旋轉（非31位）
    return x ^ ((x << 2) | (x >> 30)) ^ 
           ((x << 10) | (x >> 22)) ^ 
           ((x << 18) | (x >> 14)) ^ 
           ((x << 24) | (x >> 8));
}

/* 線性變換 L2 */
static uint32_t L2(uint32_t x) {
    return x ^ ((x << 8) | (x >> 24)) ^ 
           ((x << 14) | (x >> 18)) ^ 
           ((x << 22) | (x >> 10)) ^ 
           ((x << 30) | (x >> 2));
}

/* 非線性函數 F */
static uint32_t F() {
    uint32_t W = (X0 ^ R1) + R2;
    uint32_t W1 = R1 + X1;
    uint32_t W2 = R2 ^ X2;
    
    uint32_t u = (W1 << 16) | (W2 >> 16);
    uint32_t v = (W2 << 16) | (W1 >> 16);
    
    u = L1(u);
    v = L2(v);
    
    // 應用S盒（確保使用正確的S盒索引）
    R1 = ((uint32_t)S0[u >> 24] << 24) |
         ((uint32_t)S1[(u >> 16) & 0xFF] << 16) |
         ((uint32_t)S0[(u >> 8) & 0xFF] << 8) |
         ((uint32_t)S1[u & 0xFF]);
    
    R2 = ((uint32_t)S0[v >> 24] << 24) |
         ((uint32_t)S1[(v >> 16) & 0xFF] << 16) |
         ((uint32_t)S0[(v >> 8) & 0xFF] << 8) |
         ((uint32_t)S1[v & 0xFF]);
    
    return W;
}

/* 密鑰加載 */
void zuc_setup(const uint8_t key[16], const uint8_t iv[16]) {
    for (int i = 0; i < 16; i++) {
        lfsr[i] = ((uint32_t)key[i] << 23) | 
                  ((uint32_t)D[i] << 8) | 
                  iv[i];
    }
    
    R1 = R2 = 0;
    
    // 32輪初始化
    for (int i = 0; i < 32; i++) {
        bit_reorganize();
        uint32_t t = F();
        lfsr_init_mode(t >> 1);
    }
}

/* PRGA */
void zuc_prga(uint32_t *out, int len) {
    // 丟棄第一個輸出
    bit_reorganize();
    F();
    lfsr_work_mode();
    
    // 生成密鑰流
    for (int i = 0; i < len; i++) {
        bit_reorganize();
        out[i] = F() ^ X3;
        lfsr_work_mode();
    }
}

